package cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V2;

import cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V1.MnistTrain;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.io.ClassPathResource;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.NormalizerStandardize;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import org.nd4j.linalg.lossfunctions.LossFunctions;
import org.nd4j.linalg.api.ops.impl.transforms.comparison.Eps;
import org.nd4j.linalg.indexing.conditions.Condition;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;

/**
 * MLP、DAE版本图像解析（Chatgpt给的）
 */
public class MnistPredict {

    public static void main(String[] args) throws IOException {

        // Load the model
        MultiLayerNetwork model = ModelSerializer.restoreMultiLayerNetwork(new File(MnistTrain.BASE_PATH +"/mlp_mnist_model.zip"), true);

        // Read the test image
        BufferedImage img = ImageIO.read(new File(MnistTrain.BASE_PATH + "/mnist_png/decompression/25790.png"));

        // Preprocess the image
        INDArray image = preprocessImage(img);

        // Make predictions on the image by inputting it to the model
        INDArray output = model.output(image);

        // Get the predicted label
        int predictedLabel = Nd4j.argMax(output, 1).getInt(0);

        System.out.println("Predicted label: " + predictedLabel);
    }

    private static INDArray preprocessImage(BufferedImage image) {

        // Resize the image to 28x28
        BufferedImage resizedImage = new BufferedImage(28, 28, BufferedImage.TYPE_BYTE_GRAY);
        resizedImage.getGraphics().drawImage(image, 0, 0, 28, 28, null);

        // Convert the image to INDArray
        INDArray array = Nd4j.create(1, 784);
        int[] pixels = new int[28 * 28];
        resizedImage.getRaster().getPixels(0, 0, 28, 28, pixels);
        for (int i = 0; i < pixels.length; i++) {
            array.putScalar(i, pixels[i]);
        }

        // Normalize the image
        double mean = array.meanNumber().doubleValue();
        double std = array.stdNumber().doubleValue();
        array.subi(mean).divi(std);

        return array;
    }
}